from pydantic import BaseModel, Field
from langchain_core.messages import ToolMessage
from typing import Callable
from Agents.DialogState import State
from langgraph.prebuilt import tools_condition
from langgraph.graph import END
from tools.product_tools import product_safe_tools,product_sensitive_tools
from tools.after_sales_tools import after_sales_safe_tools,after_sales_sensitive_tools
from tools.check_orders_tools import orders_safe_tools,orders_sensitive_tools
from typing import Literal

class CompleteOrEscalate(BaseModel):
    """用于将当前任务标记为已完成或将对话框控制升级到主助手的工具，
    谁可以根据用户的需要重新路由对话框。"""


    cancel: bool = True
    reason: str

    class Config:
        json_schema_extra = {
            "example": {
                "cancel": True,
                "reason": "用户改变了对当前任务的看法。",
            },
            "example 2": {
                "cancel": True,
                "reason": "我完全完成了任务。",
            },

            "example 3": {
                "cancel": False,
                "reason": "我需要查询平台优惠政策和退换货政策。",
            },

        }

class ToProductAgent(BaseModel):
    """转移工作到一个专门的助理来处理有关商品的查询、推荐、购买等。"""

    request: str = Field(
        description="用户关于商品的任何其他信息或请求。"
    )

class ToAfterSalesAgent(BaseModel):
    """转移工作到一个专门的助理来处理有关售后服务的查询、申请、更新等。"""


    request: str = Field(
        description="用户关于售后服务的任何其他信息或请求。"
    )

class ToOrdersAgent(BaseModel):
    """转移工作到一个专门的助理来处理有关订单的查询、更新等。"""
    request: str = Field(
        description="用户关于订单的任何其他信息或请求。"
    )


def create_entry_node(assistant_name: str, new_dialog_state: str) -> Callable:
    def entry_node(state: State) -> dict:
        tool_call_id = state["messages"][-1].tool_calls[0]["id"]
        return {
            "messages": [
                ToolMessage(
                    content=f"现在您是 {assistant_name}。 反思上面主机助理和用户之间的对话。"
                    f" 用户的需求未得到满足。使用提供的工具帮助用户。记住，您是 {assistant_name}，"
                    " 只有在成功调用适当的工具后，推荐、购买或其他操作才算完成。"
                    " 如果用户改变主意或需要其他任务的帮助，请调用 CompleteOrEscalate 函数将控制权交给主要主机助理。"
                    " 不要提到您是谁 - 只需作为助理的代理。",
                    tool_call_id=tool_call_id,
                )
            ],
            "dialog_state": new_dialog_state,
        }

    return entry_node



def route_of_product_assistant(
    state: State,
):
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    did_cancel = any(tc["name"] == CompleteOrEscalate.__name__ for tc in tool_calls)
    if did_cancel:
        return "leave_node"
    safe_toolnames = [t.name for t in product_safe_tools]
    if all(tc["name"] in safe_toolnames for tc in tool_calls):
        return "product_safe_tools"
    return "product_sensitive_tools"



def route_of_after_sales_assistant(
    state: State,
):
    route = tools_condition(state)

    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    did_cancel = any(tc["name"] == CompleteOrEscalate.__name__ for tc in tool_calls)
    if did_cancel:
        return "leave_node"
    safe_toolnames = [t.name for t in after_sales_safe_tools]
    if all(tc["name"] in safe_toolnames for tc in tool_calls):
        return "after_sales_safe_tools"
    return "after_sales_sensitive_tools"


def route_of_orders_assistant(
    state: State,
):
    route = tools_condition(state)


    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    did_cancel = any(tc["name"] == CompleteOrEscalate.__name__ for tc in tool_calls)
    if did_cancel:
        return "leave_node"
    safe_toolnames = [t.name for t in orders_safe_tools]
    if all(tc["name"] in safe_toolnames for tc in tool_calls):
        return "orders_safe_tools"
    return "orders_sensitive_tools"



def pop_dialog_state(state: State) -> dict:
    """弹出对话框堆栈并返回主助手。
    这让完整的图形显式跟踪对话流并委托控制到特定的子图。
    """
    messages = []
    if state["messages"][-1].tool_calls:
        messages.append(
            ToolMessage(
                content="恢复与主机助手的对话。请反思过去的对话并根据需要协助用户。",
                tool_call_id=state["messages"][-1].tool_calls[0]["id"],
            )
        )
    return {
        "dialog_state": "pop",
        "messages": messages,
    }


#####用户可以直接和某个专业助理对话，当用户响应式，要返回当前工作流
def route_to_workflow(
    state: State,
) -> Literal[
    "primary_assistant",
    "product_assistant",
    "after_sales_assistant",
    "orders_assistant",
]:
    """如果我们处于委托状态，请直接路由到适当的助手。"""
    dialog_state = state.get("dialog_state")
    if not dialog_state:
        return "primary_assistant"
    return dialog_state[-1]


def route_primary_assistant(
    state: State,
):
    route = tools_condition(state)
    if route == END:
        return END
    tool_calls = state["messages"][-1].tool_calls
    if tool_calls:
        if tool_calls[0]["name"] == ToProductAgent.__name__:
            return "enter_product_assistant"
        elif tool_calls[0]["name"] == ToAfterSalesAgent.__name__:
            return "enter_after_sales_assistant"
        elif tool_calls[0]["name"] == ToOrdersAgent.__name__:
            return "enter_orders_assistant"
        return "primary_assistant_tools"
    raise ValueError("Invalid route")